// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <kllm/core/control_vector.h>
#include <turbo/log/logging.h>

namespace kllm {

    static ControlVectorData common_control_vector_load_one(const ControlVectorLoadInfo & load_info) {
        ControlVectorData result = { -1, {} };

        ggml_context * ctx = nullptr;
        struct gguf_init_params meta_gguf_params = {
                /* .no_alloc = */ false,
                /* .ctx      = */ &ctx,
        };
        struct gguf_context * ctx_gguf = gguf_init_from_file(load_info.fname.c_str(), meta_gguf_params);
        if (!ctx_gguf) {
            LOG(ERROR)<<__func__<<": failed to load control vector file from "<< load_info.fname.c_str();
            return result;
        }

        int32_t n_tensors = gguf_get_n_tensors(ctx_gguf);
        if (n_tensors == 0) {
            LOG(WARNING)<<__func__<<": no direction tensors found in %s"<<load_info.fname.c_str();
        }

        for (int i = 0; i < n_tensors; i++) {
            std::string name = gguf_get_tensor_name(ctx_gguf, i);

            int layer_idx = -1;

            // split on '.'
            size_t dotpos = name.find('.');
            if (dotpos != std::string::npos && name.substr(0, dotpos) == "direction") {
                try {
                    layer_idx = std::stoi(name.substr(dotpos + 1));
                } catch (...) {
                    layer_idx = -1;
                }
            }
            if (layer_idx < 0) {
                LOG(ERROR)<<__func__<<": invalid/unparsable direction tensor layer index in "<< load_info.fname.c_str();
                result.n_embd = -1;
                break;
            } else if (layer_idx == 0) {
                LOG(ERROR)<<__func__<<": invalid (zero) direction tensor layer index in "<< load_info.fname.c_str();
                result.n_embd = -1;
                break;
            }

            struct ggml_tensor * tensor = ggml_get_tensor(ctx, name.c_str());
            if (tensor->type != GGML_TYPE_F32) {
                LOG(ERROR)<<__func__<<": invalid (non-F32) direction tensor type in "<< load_info.fname.c_str();
                result.n_embd = -1;
                break;
            }
            if (ggml_n_dims(tensor) != 1) {
                LOG(ERROR)<<__func__<<": invalid (non-1D) direction tensor shape in "<< load_info.fname.c_str();
                result.n_embd = -1;
                break;
            }

            if (result.n_embd == -1) {
                result.n_embd = ggml_nelements(tensor);
            } else if (ggml_nelements(tensor) != result.n_embd) {
                LOG(ERROR)<<__func__<<": direction tensor in %s does not match previous dimensions"<< load_info.fname.c_str();
                result.n_embd = -1;
                break;
            }

            // extend if necessary - do not store data for layer 0 (it's not used)
            result.data.resize(std::max(result.data.size(), static_cast<size_t>(result.n_embd * layer_idx)), 0.0f);

            const float * src = (const float *) tensor->data;
            float * dst = result.data.data() + result.n_embd * (layer_idx - 1);  // layer 1 at [0]
            for (int j = 0; j < result.n_embd; j++) {
                dst[j] += src[j] * load_info.strength;  // allows multiple directions for same layer in same file
            }

        }

        if (result.n_embd == -1) {
            LOG(WARNING)<<__func__<<": skipping %s due to invalid direction tensors"<< load_info.fname.c_str();
            result.data.clear();
        }

        gguf_free(ctx_gguf);
        ggml_free(ctx);

        return result;
    }

    ControlVectorData common_control_vector_load(const std::vector<ControlVectorLoadInfo> & load_infos) {
        ControlVectorData result = { -1, {} };

        for (const auto & info : load_infos) {
            auto cur = common_control_vector_load_one(info);

            if (cur.n_embd == -1) {
                result.n_embd = -1;
                break;
            }
            if (result.n_embd != -1 && result.n_embd != cur.n_embd) {
                LOG(ERROR)<<__func__<<": control vectors in %s does not match previous dimensions"<< info.fname.c_str();
                result.n_embd = -1;
                break;
            }

            if (result.n_embd == -1) {
                result = std::move(cur);
            } else {
                result.data.resize(std::max(result.data.size(), cur.data.size()), 0.0f);  // extend if necessary
                for (size_t i = 0; i < cur.data.size(); i++) {
                    result.data[i] += cur.data[i];
                }
            }
        }

        if (result.n_embd == -1) {
            LOG(ERROR)<<__func__<<": no valid control vector files passed";
            result.data.clear();
        }

        return result;
    }

}